from BigDataWeb.algorithm import Algorithm
from keras.models import Sequential
from keras.layers import Dense
import numpy as np
from matplotlib import pyplot as plt
from multiprocessing import Process, Manager
from keras.layers import LSTM
from sklearn.preprocessing import MinMaxScaler
import math
from sklearn.metrics import mean_squared_error


def runLSTMTimeseriesModel(data):
    np.random.seed(1986)
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams['axes.unicode_minus'] = False
    
    lstm_units = data["lstm_units"]
    train_x = data["train_x"]
    train_y = data["train_y"]
    test_x = data["test_x"]
    test_y = data["test_y"]
    scaler = data["scaler"]
    epochs = data["epochs"]
    chart_path = data["chart_path"]
    look_back = data["look_back"]
    dataset = data["dataset"]
    
    # 构建模型
    model = Sequential()
    model.add(LSTM(units=lstm_units, input_shape=(1, look_back)))
    model.add(Dense(1))
    model.compile(loss='mean_squared_error', optimizer='adam')

    # 打印模型信息
    summary_info = []
    model.summary(print_fn=lambda x:summary_info.append(x))
    data["model_summary"] = "<br/>".join(summary_info)
    
    # 编译并训练模型
    history = model.fit(train_x, train_y, epochs=epochs, batch_size=1, verbose=0)
    
    # 绘制训练过程(损失函数)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(history.history["loss"])
    ax.set_title("损失函数")
    ax.set_xlabel("epoch")
    ax.set_ylabel("loss_value")
    fig.savefig("%s/loss_value.png" % (chart_path))
    
    # 训练集
    train_predict = model.predict(train_x)
    train_predict = scaler.inverse_transform(train_predict)
    train_y = scaler.inverse_transform(train_y.reshape(-1, 1))
    train_score = math.sqrt(mean_squared_error(train_predict, train_y))
    data["train_score"] = train_score
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_title("训练集")
    ax.plot(train_y)
    ax.plot(train_predict)
    fig.legend(["真实值", "预测值"])
    fig.savefig("%s/train.png" % (chart_path))
    
    # 测试集
    test_predict = model.predict(test_x)
    test_predict = scaler.inverse_transform(test_predict)
    test_y = scaler.inverse_transform(test_y.reshape(-1, 1))
    test_score = math.sqrt(mean_squared_error(test_predict, test_y))
    data["test_score"] = test_score
    
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_title("测试集")
    ax.plot(test_y)
    ax.plot(test_predict)
    fig.legend(["真实值", "预测值"])
    fig.savefig("%s/test.png" % (chart_path))
    
    # 预测下一个周期的值
    input_x = dataset[-1 * look_back:, :]
    input_x = input_x.reshape(input_x.shape[1], 1, input_x.shape[0])
    output_y = model.predict(input_x)
    output_y = scaler.inverse_transform(output_y.reshape(-1, 1))
    output_y = output_y[0][0]
    data["output_y"] = output_y


class LSTMTimeseries(Algorithm):
    in_out_put_filed_name = ""
    look_back = 1
    train_ratio = 0.8
    lstm_units = 8
    epochs = 100
    
    scaler = None
    raw_dataset = None
    dataset = None
    train_x = None
    train_y = None
    test_x = None
    test_y = None
    
    model_summary = ""
    train_score = 0
    test_score = 0
    output_y = 0
    
    def prepareIpynbItems(self):
        Algorithm.prepareIpynbItems(self)
        self.ipynb_items["#in_out_put_filed_name#"] = self.in_out_put_filed_name
        self.ipynb_items["#look_back#"] = self.look_back
        self.ipynb_items["#train_ratio#"] = self.train_ratio
        self.ipynb_items["#lstm_units#"] = self.lstm_units
        self.ipynb_items["#epochs#"] = self.epochs

    def __init__(self):
        Algorithm.__init__(self)
        self.algorithm_name = "深度学习-LSTM"
        self.ipynb_template_name = "lstm_timeseries-template.ipynb"
        
    # 转换数据集
    def createDataset(self, dataset, look_back=1):
        dataX, dataY = [], []
        for i in range(len(dataset) - look_back):
            a = dataset[i:(i + look_back), 0]
            dataX.append(a)
            dataY.append(dataset[i + look_back, 0])
        return np.array(dataX), np.array(dataY)
    
    def setInOutPutFieldName(self, in_out_put_filed_name):
        self.in_out_put_filed_name = in_out_put_filed_name
        self.raw_dataset = self.data_source[[in_out_put_filed_name]].values
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        self.dataset = self.scaler.fit_transform(self.raw_dataset)

    def implent(self):
        train_size = int(len(self.dataset) * self.train_ratio)
        train = self.dataset[0:train_size, :]
        test = self.dataset[train_size:len(self.dataset), :]
        self.train_x, self.train_y = self.createDataset(train, self.look_back)
        self.test_x, self.test_y = self.createDataset(test, self.look_back)
        self.train_x = np.reshape(self.train_x, (self.train_x.shape[0], 1, self.train_x.shape[1]))
        self.test_x = np.reshape(self.test_x, (self.test_x.shape[0], 1, self.test_x.shape[1]))
        
        with Manager() as manager:
            data = manager.dict()
            data["lstm_units"] = self.lstm_units
            data["train_x"] = self.train_x
            data["train_y"] = self.train_y
            data["test_x"] = self.test_x
            data["test_y"] = self.test_y
            data["scaler"] = self.scaler
            data["epochs"] = self.epochs
            data["chart_path"] = self.chart_path
            data["look_back"] = self.look_back
            data["dataset"] = self.dataset
            
            process = Process(target=runLSTMTimeseriesModel, args=(data,))
            process.start()
            process.join()
            
            self.model_summary = data["model_summary"]
            self.train_score = data["train_score"]
            self.test_score = data["test_score"]
            self.output_y = data["output_y"]
